# -*- coding:utf-8 -*-
"""
@filename   : lesson-04-Computational-Graph
@author     : QuZhang
@date       : 2020-11-22
@brief      : 计算图示例
"""
import torch


if __name__ == "__main__":
    # y = (w+x) * (w+1)
    # 计算图：函数里的自变量是叶子节点，叶子节点是需要学习的参数
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    a = torch.add(w, x)
    a.retain_grad()  # 设置报存非叶子节点的梯度
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    y.backward()
    print("w.grad: ", w.grad)

    # 查看叶子节点
    print("is_leaf:\n", w.is_leaf, x.is_leaf, a.is_leaf, b.is_leaf, y.is_leaf)

    # 查看梯度
    # 节约内存考虑，最后不保存非叶子节点的梯度，只保存叶子节点梯度
    print("gradient:\n", w.grad, x.grad, a.grad, b.grad, y.grad)

    # 查看 grad_fn
    # 记录相应的变量是由什么操作得到的
    print("grad_fn:\n", w.grad_fn, x.grad_fn, a.grad_fn, b.grad_fn, y.grad_fn)
